import argparse
import os
import numpy as np
import torch
import yaml
from tqdm import tqdm
import archs
import albumentations as A
from dataset import TestDataset
from utils import str2bool
import cv2  # 添加cv2导入


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--name', required=True, help='model name')
    parser.add_argument('--test_dataset', default='Test/Rural', help='test dataset name')
    parser.add_argument('--batch_size', default=4, type=int, help='batch size')
    parser.add_argument('--num_workers', default=4, type=int, help='number of workers')
    parser.add_argument('--output_dir', default='results', help='output directory')
    parser.add_argument('--save_probability', default=False, type=str2bool, help='save probability maps')
    parser.add_argument('--model_path', default=None, type=str, help='path to model checkpoint')
    return parser.parse_args()


def main():
    config = vars(parse_args())

    # 创建输出目录
    os.makedirs(os.path.join(config['output_dir'], config['name']), exist_ok=True)
    os.makedirs(os.path.join(config['output_dir'], config['name'], 'predictions'), exist_ok=True)
    if config['save_probability']:
        os.makedirs(os.path.join(config['output_dir'], config['name'], 'probabilities'), exist_ok=True)

    # 加载模型配置
    with open('models/%s/config.yml' % config['name'], 'r') as f:
        model_config = yaml.load(f, Loader=yaml.FullLoader)

    # 创建模型
    print("=> creating model %s" % model_config['arch'])
    model = archs.__dict__[model_config['arch']](model_config['num_classes'],
                                                 model_config['input_channels'],
                                                 model_config['deep_supervision'])
    model = model.cuda()

    # 加载模型权重（修复安全警告）
    model_path = config['model_path'] if config['model_path'] else 'models/%s/model.pth' % config['name']
    print('=> loading model %s' % model_path)
    model.load_state_dict(torch.load(model_path, weights_only=True))  # 添加 weights_only=True
    model.eval()

    # 创建测试数据集
    test_transform = A.Compose([
        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ])

    # 构建正确的数据集路径（修改为 images_png）
    test_img_dir = os.path.join('inputs', config['test_dataset'], 'images_png')

    # 检查路径是否存在
    if not os.path.exists(test_img_dir):
        print(f"警告: 图像目录不存在: {test_img_dir}")
        # 尝试使用原始路径
        test_img_dir = os.path.join('inputs', config['test_dataset'], 'images')
        if not os.path.exists(test_img_dir):
            raise ValueError(f"图像目录不存在: {test_img_dir}")

    # 获取图像ID列表
    img_ids = [os.path.splitext(f)[0] for f in os.listdir(test_img_dir) if f.endswith(model_config['img_ext'])]

    test_dataset = TestDataset(
        img_ids=img_ids,
        img_dir=test_img_dir,
        img_ext=model_config['img_ext'],
        transform=test_transform,
        target_size=(model_config['input_h'], model_config['input_w'])
    )

    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=config['batch_size'],
        shuffle=False,
        num_workers=config['num_workers'],
        drop_last=False)

    # 进行预测
    with torch.no_grad():
        for input, img_ids in tqdm(test_loader, total=len(test_loader)):
            input = input.cuda()

            # 模型预测
            output = model(input)

            if model_config['deep_supervision']:
                output = output[-1]

            # 保存预测结果
            output = torch.softmax(output, dim=1)
            predictions = torch.argmax(output, dim=1).cpu().numpy()

            for i in range(len(img_ids)):
                img_id = img_ids[i]
                pred = predictions[i]

                # 保存预测图像
                cv2.imwrite(
                    os.path.join(config['output_dir'], config['name'], 'predictions',
                                 img_id + model_config['mask_ext']),
                    pred.astype(np.uint8)
                )

                # 保存概率图
                if config['save_probability']:
                    prob = output[i].cpu().numpy()
                    np.save(
                        os.path.join(config['output_dir'], config['name'], 'probabilities', img_id + '.npy'),
                        prob
                    )


if __name__ == '__main__':
    main()